#!/usr/bin/env python
# -*- coding: utf-8 -*-

import socket
import threading
import time
from typing import Optional, Tuple, Dict
from logger import logger
from utils import parse_address
from config import config

class NodeDiscovery:
    def __init__(self):
        self.stun_server = parse_address(config.get('stun_server'))
        self.local_address = None
        self.public_address = None
        self.nat_type = None
        self.remote_address = None
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind(('0.0.0.0', 0))
        self.running = False
        self.listen_thread = None

    def start(self) -> None:
        """启动节点发现服务"""
        self.running = True
        self.local_address = self.socket.getsockname()
        self.detect_nat()
        self.listen_thread = threading.Thread(target=self._listen)
        self.listen_thread.daemon = True
        self.listen_thread.start()
        logger.info(f'节点发现服务启动，本地地址: {self.local_address}')

    def stop(self) -> None:
        """停止节点发现服务"""
        self.running = False
        if self.socket:
            self.socket.close()
        if self.listen_thread:
            self.listen_thread.join()
        logger.info('节点发现服务停止')

    def detect_nat(self) -> None:
        """检测NAT类型和公网地址"""
        try:
            # 发送探测请求到STUN服务器
            self.socket.sendto(b'\x00\x01\x00\x00', self.stun_server)
            data, addr = self.socket.recvfrom(1024)
            if data:
                self.public_address = addr
                logger.info(f'获取到公网地址: {self.public_address}')
        except Exception as e:
            logger.error(f'NAT检测失败: {e}')

    def send_discovery_message(self, remote_address: Tuple[str, int]) -> None:
        """发送发现消息到远程节点"""
        try:
            message = f'DISCOVER:{config.get("node_id")}'.encode()
            self.socket.sendto(message, remote_address)
            logger.debug(f'发送发现消息到 {remote_address}')
        except Exception as e:
            logger.error(f'发送发现消息失败: {e}')

    def _listen(self) -> None:
        """监听来自其他节点的消息"""
        while self.running:
            try:
                data, addr = self.socket.recvfrom(1024)
                if data:
                    message = data.decode()
                    if message.startswith('DISCOVER:'):
                        remote_node_id = message.split(':')[1]
                        self.handle_discovery(remote_node_id, addr)
            except Exception as e:
                if self.running:
                    logger.error(f'监听消息失败: {e}')

    def handle_discovery(self, remote_node_id: str, addr: Tuple[str, int]) -> None:
        """处理发现消息"""
        logger.info(f'收到来自节点 {remote_node_id} 的发现消息，地址: {addr}')
        self.remote_address = addr
        # 发送响应消息
        response = f'DISCOVER_ACK:{config.get("node_id")}'.encode()
        self.socket.sendto(response, addr)

    def get_remote_address(self) -> Optional[Tuple[str, int]]:
        """获取远程节点地址"""
        return self.remote_address